package com.senke.aiagent.tools;

import cn.hutool.http.ContentType;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;

import java.util.List;
import java.util.stream.Collectors;

/**
 * 搜索工具，调用智谱 AI 搜索引擎的 API
 */
@Slf4j
public class ZhiPuAISearchTool {

    // SearchAPI 的搜索接口地址
    private static final String SEARCH_API_URL = "https://open.bigmodel.cn/api/paas/v4/web_search";

    private final String apiKey;

    public ZhiPuAISearchTool(String apiKey) {
        this.apiKey = apiKey;
    }

    @Tool(description = "Search for information from Baidu Search Engine")
    public String searchWeb(@ToolParam(description = "Search query keyword") String query) {
        // 1. 构建请求参数
        JSONObject requestBody = new JSONObject();
        requestBody.set("search_engine", "search_std")  // 必填参数
                .set("search_query", query);            // 必填参数
        try {
            // 2. 发送POST请求
            HttpResponse response = HttpRequest.post(SEARCH_API_URL)
                    .header("Authorization", "Bearer " + apiKey) // 认证头
                    .contentType(ContentType.JSON.getValue())     // JSON格式
                    .body(requestBody.toString())                  // 请求体
                    .timeout(5000)                       // 超时时间
                    .execute();
            // 3. 处理响应
            if (response.getStatus() == 200) {
                JSONObject responseBody = JSONUtil.parseObj(response.body());
                // 解析搜索结果
                JSONArray searchResults = responseBody.getJSONArray("search_result");
                List<Object> objects = searchResults.subList(0, Math.min(5, searchResults.size()));
                // 拼接搜索结果为字符串
                String result = objects.stream().map(obj -> {
                    JSONObject newJSONObject = new JSONObject();
                    JSONObject tmpJSONObject = (JSONObject) obj;
                    newJSONObject.set("title", tmpJSONObject.get("title"))
                            .set("content", tmpJSONObject.get("content"));
                    return newJSONObject.toString();
                }).collect(Collectors.joining(","));
                return result;
            } else {
                return "Error searching: " + response.body();
            }
        } catch (Exception e) {
            return "Error searching: " + e.getMessage();
        }
    }
}